from mmcv.cnn import ConvModule, DepthwiseSeparableConvModule

class UNetDecodeHead(nn.Module):
    def __init__(self,
                 in_channels,
                 channels,
                 in_index,
                 num_classes,
                 kernel_size=3,
                 concat_input=True,
                 dropout_ratio=-1,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU')):
        super(UNetDecodeHead, self).__init__()

        self.in_channels = in_channels
        self.channels = channels
        self.in_index = in_index
        self.num_classes = num_classes
        self.kernel_size = kernel_size
        self.concat_input = concat_input
        self.dropout_ratio = dropout_ratio
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg

        self.conv1 = ConvModule(
            self.in_channels[-1],
            self.channels,
            kernel_size=self.kernel_size,
            padding=self.kernel_size // 2,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)

        self.conv2 = ConvModule(
            self.channels,
            self.channels,
            kernel_size=self.kernel_size,
            padding=self.kernel_size // 2,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)

        self.conv3 = ConvModule(
            self.channels * 2 if self.concat_input else self.channels,
            self.channels,
            kernel_size=self.kernel_size,
            padding=self.kernel_size // 2,
            conv_cfg=self.conv_cfg,
            norm_cfg=self.norm_cfg,
            act_cfg=self.act_cfg)

        self.conv4 = ConvModule(
            self.channels,
            self.num_classes,
            kernel_size=1,
            conv_cfg=self.conv_cfg,
            act_cfg=None)

    def forward(self, inputs):
        x = self.conv1(inputs[-1])
        x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)

        if self.concat_input:
            x = torch.cat([x, inputs[-2]], dim=1)

        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)

        return x
